import torch
import torch.nn as nn
import torch.nn.functional as F

class CBAM(nn.Module):
    """Convolutional Block Attention Module"""
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, channels // reduction, 1),
            nn.ReLU(),
            nn.Conv2d(channels // reduction, channels, 1),
            nn.Sigmoid()
        )
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(2, 1, 7, padding=3),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        # Channel attention
        ca = self.channel_attention(x)
        x = x * ca
        
        # Spatial attention
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        spatial_input = torch.cat([avg_out, max_out], dim=1)
        sa = self.spatial_attention(spatial_input)
        x = x * sa
        
        return x

class YOLOv7Pose(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.backbone = self._build_backbone()
        self.neck = self._build_neck()
        self.head = self._build_head()
        
    def _build_backbone(self):
        """构建带CBAM的CSPDarknet53 backbone"""
        # 简化的backbone结构
        layers = []
        in_channels = self.config['input_channels']
        
        # 初始卷积层
        layers.append(nn.Conv2d(in_channels, 32, 3, padding=1))
        layers.append(nn.BatchNorm2d(32))
        layers.append(nn.SiLU(inplace=True))
        
        # CSPDarknet53结构（简化）
        channels = [32, 64, 128, 256, 512]
        for i in range(len(channels)-1):
            layers.append(self._make_csp_layer(channels[i], channels[i+1]))
            layers.append(CBAM(channels[i+1]))
        
        return nn.Sequential(*layers)
    
    def _make_csp_layer(self, in_channels, out_channels):
        """构建CSP层"""
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True)
        )
    
    def _build_neck(self):
        """构建BiFPN颈部网络"""
        # 简化的BiFPN结构
        return nn.Sequential(
            nn.Conv2d(512, 256, 1),
            nn.BatchNorm2d(256),
            nn.SiLU(inplace=True)
        )
    
    def _build_head(self):
        """构建检测头"""
        return nn.ModuleDict({
            'bbox': nn.Conv2d(256, 4, 1),  # 边界框
            'conf': nn.Conv2d(256, 1, 1),  # 置信度
            'keypoints': nn.Conv2d(256, self.config['num_keypoints'] * 3, 1)  # 关键点
        })
    
    def forward(self, x):
        """
        Args:
            x: [batch_size, seq_len, 1, H, W]
        Returns:
            detections: [batch_size, seq_len, max_objects, 6 + 17*3]
        """
        batch_size, seq_len = x.shape[:2]
        x = x.reshape(batch_size * seq_len, *x.shape[2:])
        
        # Backbone特征提取
        features = self.backbone(x)
        
        # Neck特征融合
        neck_features = self.neck(features)
        
        # Head预测
        bbox_pred = self.head['bbox'](neck_features)
        conf_pred = self.head['conf'](neck_features)
        kp_pred = self.head['keypoints'](neck_features)
        
        # 后处理（简化）
        detections = self._post_process(bbox_pred, conf_pred, kp_pred, batch_size, seq_len)
        
        return detections
    
    def _post_process(self, bbox_pred, conf_pred, kp_pred, batch_size, seq_len):
        """后处理得到最终检测结果"""
        # 这里实现NMS等后处理逻辑
        # 返回格式: [batch_size, seq_len, max_objects, 6 + 51]
        # 6: [conf, x, y, w, h, class], 51: 17个关键点 * 3 (x, y, visibility)
        max_objects = 50  # 最大检测对象数
        
        # 简化的后处理
        detections = torch.zeros(batch_size, seq_len, max_objects, 6 + 51)
        
        return detections